'''
Created on May 10, 2011

@author: oabalbin
'''

import exome.vcftools.SNP_positions as pos 
from random import sample
from optparse import OptionParser
from collections import defaultdict

def randomize_vcf(array_size, db_frac):
    '''
    Returns an array of random indexes
    '''
    vec=range(array_size)
    return sample(vec,int(round(db_frac*array_size)))

def rand_vect_dict(rand_vec):
    '''
    '''
    rand_dict = defaultdict()
    for i in rand_vec:
        rand_dict[i]=i 
    return rand_dict

def read_file(input_file_name,out_file_name, random_vector):
    '''
    if the number line is in the array of randomly generated
    integers print the lines otherwise skkip it    
    '''
    input_file = open(input_file_name)
    outfile=open(out_file_name,'w')
    i=0
    print 'Randomizing %s'%(input_file_name)
    for l in input_file:
        if l.startswith('#'):
            outfile.write(l)
            continue
        try:
            random_vector[i]
        #if i in random_vector:
            outfile.write(l)
        except KeyError:
            pass
        i+=1
       
    outfile.close()


def make_pos_file(ifile, ofile):
    '''
    '''
    ifile = open(ifile)
    ofile =open(ofile,'w')
    
    for l in ifile:
        if l.startswith('#'):
            ofile.write(l)
            continue
        fields=l.strip('\n').replace(',','$').split('\t')
        
        pos= '%s@%s'%(fields[0],fields[1])
        ofile.write(pos+',')
    
    ifile.close()
    ofile.close()



def main_rand_snp(input_file_name, input_pos, out_file_name):
    '''
    It writes a random dbSNP that is a fraction (db_frac) of the
    original one.
    '''
    db_frac=0.25
    array_size = len(pos.read_positions(input_pos))
    rand_vec = randomize_vcf(array_size,db_frac)
    rand_dict = rand_vect_dict(rand_vec)
    read_file(input_file_name,out_file_name, rand_dict)


if __name__ == '__main__':
        
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--query_vcf_file", dest="query_vcf_file",
                            help="query_vcf_file")
    optionparser.add_option("-p", "--query_pos_file", dest="query_pos_file",
                            help="query_pos_file")


    (options, args) = optionparser.parse_args()     
    outfile=options.query_vcf_file.replace('.vcf','_random.vcf')
    if not options.query_pos_file:
        make_pos_file(options.query_vcf_file,options.query_vcf_file.replace('.vcf','.pos'))
        query_pos_file=options.query_vcf_file.replace('.vcf','.pos')
    else:
        query_pos_file=options.query_pos_file
        
    main_rand_snp(options.query_vcf_file, query_pos_file, outfile)

    
    